import os.path

import pandas as pd
from src.EnvironmentVariables import BASE_PATH


def read_file(name: str, is_res=False, use_default_arg=True, **kwargs):
    if use_default_arg:
        if 'parse_dates' not in kwargs.keys():
            kwargs['parse_dates'] = [0]
        if 'index_col' not in kwargs.keys():
            kwargs['index_col'] = [0]
    tar_path = os.path.join(BASE_PATH, 'data/indexDataDay/' + ("Preprossed" if is_res else ""), name)
    if 'test.csv' in name:
        return pd.read_csv(tar_path, **kwargs)
    return pd.read_excel(tar_path, **kwargs)


def make_path(name: str, is_res=False):
    return os.path.join(BASE_PATH, 'data/indexDataDay/' + ("Preprossed" if is_res else ""), name)


def dump_opened_date():
    df = read_file("指数行情序列（交易特征）.xlsx", header=[0, 1], parse_dates=[0], index_col=[0])
    date_df = df.index
    date_df = pd.DataFrame(date_df)
    date_df.rename(columns={0: 'Date'}, inplace=True)
    # 保存为CSV文件
    date_df.to_csv(make_path('OpenDate.csv', True), index=False)


def use_before_vla_fill_na():
    tar_list = ['中国_人民币存款准备金率_大型存款类金融机构(变动公告日期)', '银行间质押式回购加权利率_7天',
                '中国大宗商品价格指数_总指数', '中国_融资融券余额', '中证行业指数_能源']
    date_df = read_file('OpenDate.csv', True, parse_dates=[0], index_col=[0])
    for tar_name in tar_list:
        tar_df = read_file(tar_name + '.xlsx', parse_dates=[0], index_col=[0])
        tar_df = pd.concat([date_df, tar_df], axis=1)
        tar_df.fillna(method='ffill', inplace=True)
        tar_df = tar_df.loc[tar_df.index.isin(date_df.index)]
        tar_df.index = tar_df.index.set_names(['Date'])
        tar_df.to_csv(make_path(tar_name + 'test.csv', True))


def fill_monthly_data():
    tar_list = (
        {
            'name': 'CPI',
            'update': 11,
        },
        {
            'name': 'PPI',
            'update': 11,
        },
        {
            'name': 'PMI',
            'update': -1,
        },
        {
            'name': 'M1',
            'update': 15,
        },
        {
            'name': 'M2',
            'update': 15,
        }
    )
    date_df = read_file('OpenDate.csv', True, parse_dates=[0], index_col=[0])
    names = ['中国_CPI_当月同比', '中国_M1_同比']
    for name in names:
        df = read_file(name + '.xlsx', parse_dates=[0], index_col=[0])
        res = [date_df]
        for item in tar_list:
            ele, update = item.values()
            single_df = df.filter(like=ele)
            if update == -1:
                res.append(single_df)
            else:
                new_index = single_df.index + pd.DateOffset(months=1) - pd.offsets.MonthBegin()
                single_df.index = pd.to_datetime(new_index.strftime(f'%Y-%m-{update}'))
                res.append(single_df)
        res = pd.concat(res, axis=1)
        res.fillna(method='ffill', inplace=True)
        res = res.loc[res.index.isin(date_df.index)]
        res.to_csv(make_path(name + 'test.csv', True))


def dump_market_price():
    file_names = ['指数行情序列（交易特征）', '指数行情序列（财务特征）']
    for name in file_names:
        df = read_file(name + '.xlsx', header=[0, 1], parse_dates=[0], index_col=[0])
        df.to_csv(make_path(name + 'test.csv', True))


def split_index_data():
    def get_single_data(origin: pd.DataFrame, name: str):
        new_df = origin.loc[:, origin.columns.get_level_values(1) == name]
        new_df.columns = new_df.columns.droplevel(level=1)
        return new_df

    date_df = read_file('OpenDate.csv', True, parse_dates=[0], index_col=[0])
    file_name = ('指数行情序列（交易特征）', '指数行情序列（财务特征）', 'index_factors')
    index_name = ['中证500', '上证指数', '上证50', '沪深300']
    label_name = ['涨跌幅', '换手率']  # ['收盘价', '最高价', '最低价']
    origin_dfs = [read_file(_ + 'test.csv', True, header=[0, 1], parse_dates=[0], index_col=[0]) for _ in file_name]
    for tar_index in index_name:
        res = []
        for odf in origin_dfs:
            res.append(get_single_data(odf, tar_index))
        res = pd.concat(res, axis=1)
        res.dropna(inplace=True)
        res = res.iloc[res.index.isin(date_df.index)]
        ma5 = res.filter(like='MA简单移动平均5日')
        sign_label = date_df.copy()
        sign_label.insert(0, "Sign", 0)
        for i in range(len(ma5) - 5):
            if ma5.iloc[i + 5, 0] > ma5.iloc[i, 0]:
                sign_label.iloc[i, 0] = 1
            else:
                sign_label.iloc[i, 0] = 0
        # ma5 = ma5.shift(-5, axis=0)
        label = [res.filter(like=_) for _ in label_name]
        # label.append(ma5)
        label.append(sign_label)
        label = pd.concat(label, axis=1)
        label.dropna(inplace=True)
        label.to_csv(make_path(tar_index + '_label.csv', True))
        res = res.loc[res.index.isin(label.index)]
        res.to_csv(make_path(tar_index + 'test.csv', True))






if __name__ == '__main__':
    pass
    split_index_data()
